from kmeans_search import kmeans_search
from sklearn.linear_model import LogisticRegression


class PrecisionModel:
    def __init__(self, dataset, label, search_model=kmeans_search):
        self.dataset = dataset
        self.label = label
        self.search_model = search_model

    def fit(self, target):
        X_train, y_train = self.search_model(self.dataset, self.label, target)
        if sum(y_train) == 0:
            return np.array([1, 0])
        if sum(y_train) == len(y_train):
            return np.array([0, 1])
        lr = LogisticRegression(max_iter=1000).fit(X_train, y_train)
        return lr.predict_proba(target.reshape(1, -1))[0]

    def predict_proba(self, X_test):
        y_pred = []
        for target in X_test:
            y_pred.append(self.fit(target))
        return np.array(y_pred)

    def predict(self, X_test):
        return np.where(self.predict_proba(X_test)[:, 1] > 0.5, 1, 0)


if __name__ == '__main__':
    pm = PrecisionModel(1, 2)
